from openai import OpenAI
import os

def chatNow(model='gpt-3.5-turbo-instruct',mode='balance',prompt=None):
    #根据mode来调整temperature及presence_penalty

    if mode == 'balance':
        temperature = 1
        presence_penalty = 0
    elif mode == 'precision':
        temperature = 0.8
        presence_penalty = 2
    elif mode == 'creativity':
        temperature = 1.2
        presence_penalty = -1

    def chat(prompt):
        try:
            client = OpenAI(api_key="sk-yrpeg76LZt53LhyTcue8ojFcvRkzs53fapexPklNqR5KPKae", base_url="https://api.fe8.cn/v1")
            response = client.completions.create(model=model, prompt=prompt, max_tokens=1000, temperature=temperature, presence_penalty=presence_penalty)

            return response.choices[0].text.strip()
        except Exception as e:
            return "broken"

    print(chat(prompt))

#Zero-Shot-CoT
#chatNow(mode="balance",prompt="罗杰有五个网球，他又买了两盒网球，每盒有3个网球，请问他现在总共有多少个网球？请一步步进行推理并得出结论。")
#chatNow(mode="balance",prompt="食堂总共有23个苹果，如果他们用掉20个苹果，然后又买了6个苹果，请问现在食堂总共有多少个苹果？请一步步进行推理并得出结论。")
#chatNow(mode="balance",prompt="杂耍者可以杂耍16个球。其中一半的球是高尔夫球，其中一半的高尔夫球是蓝色的。请问总共有多少个蓝色高尔夫球？请一步步进行推理并得出结论。")
#chatNow(mode="balance",prompt="艾米需要4分钟才能爬到滑梯顶部，她花了1分钟才滑下来，水滑梯将在15分钟后关闭，请问在关闭之前她能滑多少次？请一步步进行推理并得出结论。")

#Few-Shot-CoT
# prompt= """
# Q："罗杰有五个网球，他又买了两盒网球，每盒有3个网球，请问他现在总共有多少个网球？
# A："罗杰一开始有五个网球，又购买了两盒网球，每盒3个，共购买了6个网球，因此现在总共由5+6=11个苹果，然后又买了6个苹果，请问现在食堂总共有多少个苹果？个网球。因此答案是11."
# Q:“食堂总共有23个苹果，如果他们用掉20
# A："食堂最初有23个苹果，用掉20个，然后又买了6个，总共有23-20+6=9个苹果，答案是9
# Q：“杂耍者可以杂耍16个球。其中一半的球是高尔夫球，其中一半的高尔夫球是蓝色的。请问总共有多少个蓝色高尔夫球？"
# A:“总共有16个球，其中一半是高尔夫球，也就是8个，其中一半是蓝色的，也就是4个，答案是4个。"
# Q:“艾米需要4分钟才能爬到滑梯顶部，她花了1分钟才滑下来，水滑梯将在15分钟后关闭，请问在关闭之前她能滑多少次？
# A：
# """
# chatNow(mode="balance",prompt=prompt)

chatNow(mode="balance",model= "gpt-3.5-turbo-instruct",prompt="""Q:“艾米需要4分钟才能爬到滑梯顶部，她花了1分钟才滑下来，水滑梯将在15分钟后关闭，请问在关闭之前她能滑多少次？"
A：为了解决“在关闭之前她能滑多少次？”这个问题，我们首先要解决的问题是""")